"""Parse phylogenetic trees entered in Comapct Newick Format.  CNF is a subset of Newick
(New Hampshire) format that is simpler and more fully specified.  The goal is to avoid some of the
corner cases in the Newicki format while retaining human readability and backwards compatability.

Differences between CNF and Newick:
* no square bracket comments
* only alphanumeric characters (ASCII; A-Z, a-z, 0-9) and _ are allowed in taxon names
* no spaces in taxon names
* no quoted taxon names
* the first character of a taxon name must be an alpha (A-Z, a-z)
* distances can be any (double precision) floating point number
"""

import pyparsing as pp

# Enable packrat parsing
pp.ParserElement._parse = pp.ParserElement._parseCache

# CNF Grammar
name = pp.Word(pp.alphanums + '_')
distance = pp.Suppress(":") + pp.Word(pp.nums + '.+-eE')
distance.setParseAction(lambda t: float(t[0]))
node = pp.Forward()
children = pp.Suppress("(") + pp.delimitedList(node) + pp.Suppress(")")

nodeC = children("children") + pp.Optional(name)("name") + pp.Optional(distance)("distance")
nodeL = pp.Optional(children)("children") + name("name") + pp.Optional(distance)("distance")
nodeD = pp.Optional(children)("children") + pp.Optional(name)("name") + distance("distance")
node << pp.Group(nodeC | nodeL | nodeD)

tree = node + pp.Optional(pp.Suppress(";")) + pp.stringEnd

class CNFParseException(Exception):
    """Exception raised if there is a problem with parsing a CNF string."""
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

def precheck(cnfString):
    """Do some basic sanity checking of CNF strings.  At this point, just looks for characters that 
    should not be present.  Return False if a problem was detected."""
    valid = pp.alphas + '_' + pp.nums + '.+-eE' + ',():;'
    for c in cnfString:
        if valid.find(c) == -1:
            raise CNFParseException("precheck: " + c + " is an invalid character.")
    return True

def preprocess(cnfString):
    """Handles the case where a tree is provided which is missing leaf names, for instance ((,,),), and
    insert some dummy leaf names so that the string can be parsed.  Does not check for duplicate names!"""
    cnfString = cnfString.replace("(,","(%,").replace(",)",",%)").replace("),",")%,").replace(",,",",%,")
    s = []
    i = 1
    for c in cnfString:
        if c == "%":
            s.append(str(i))
            i += 1
        else:
            s.append(c)
    return "".join(s)

def parse(cnfString):
    precheck(cnfString)
    cnfString = preprocess(cnfString)
    try:
        t = tree.parseString(cnfString)
    except pp.ParseException as e:
        raise CNFParseException("parse: " + str(e))
    return t
